﻿namespace StyleCop.CSharp
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using QuickGraph;

    /// <summary>
    /// Builds code flow graphs.
    /// </summary>
    public static class FlowGraph
    {
        private static readonly List<CodeUnitBuilder> builders = FlowGraph.BuildBuilderList();

        private delegate IEnumerable<ICodeUnit> CodeUnitBuilder(CodeUnitGraph graph, ICodeUnit codeUnit);

        #region Public Methods
        /// <summary>
        /// Builds the flow graph for the code unit.
        /// </summary>
        /// <param name="codeUnit">The code unit.</param>
        /// <returns>The flow graph for the code unit.</returns>
        public static CodeUnitGraph BuildFlowGraph(ICodeUnit codeUnit)
        {
            if (codeUnit == null)
            {
                throw new ArgumentNullException("codeUnit");
            }

            var graph = new CodeUnitGraph();

            graph.AddVertex(codeUnit);

            graph.BuildFlowGraph(Enumerable.Empty<ICodeUnit>(), codeUnit);

            return graph;
        }
        #endregion Public Methods

        #region Private Methods
        private static IEnumerable<ICodeUnit> BuildAnonymousMethodExpressionFlowGraph(CodeUnitGraph graph, AnonymousMethodExpression anonymousMethod)
        {
            var verticesFromAnonymousMethod = graph.BuildFlowGraph(
                new ICodeUnit[] { anonymousMethod },
                anonymousMethod.Children());

            return verticesFromAnonymousMethod.Concat(new ICodeUnit[] { anonymousMethod });
        }

        private static IEnumerable<ICodeUnit> BuildAssignmentExpressionFlowGraph(CodeUnitGraph graph, AssignmentExpression assignment)
        {
            /*  The assignment statements should be built into the graph with the right hand side first so that the
                variable checking rules get the accessing first before the assigning. This was an issue in a statement
                like s = s + 1; */

            var verticesFromLeftHandSide = graph.BuildFlowGraph(
                new ICodeUnit[] { assignment },
                assignment.RightHandSide);

            return graph.BuildFlowGraph(
                verticesFromLeftHandSide,
                assignment.LeftHandSide);
        }

        private static List<CodeUnitBuilder> BuildBuilderList()
        {
            var builders = new List<CodeUnitBuilder>();
            builders.Add(FlowGraph.CreateBuilder<AnonymousMethodExpression>(FlowGraph.BuildAnonymousMethodExpressionFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<AssignmentExpression>(FlowGraph.BuildAssignmentExpressionFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<CatchStatement>(FlowGraph.BuildCatchFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<DoWhileStatement>(FlowGraph.BuildDoWhileFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<ElseStatement>(FlowGraph.BuildElseFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<ForStatement>(FlowGraph.BuildForFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<ForeachStatement>(FlowGraph.BuildForeachFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<IfStatement>(FlowGraph.BuildIfFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<LambdaExpression>(FlowGraph.BuildLambdaExpressionGraph));
            builders.Add(FlowGraph.CreateBuilder<ReturnStatement>(FlowGraph.BuildReturnFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<SwitchStatement>(FlowGraph.BuildSwitchFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<ThrowStatement>(FlowGraph.BuildThrowStatementGraph));
            builders.Add(FlowGraph.CreateBuilder<TryStatement>(FlowGraph.BuildTryFlowGraph));
            builders.Add(FlowGraph.CreateBuilder<WhileStatement>(FlowGraph.BuildWhileFlowGraph));

            return builders;
        }

        private static IEnumerable<ICodeUnit> BuildCatchFlowGraph(CodeUnitGraph graph, CatchStatement catchStatement)
        {
            // This is used three times so use the same instance.
            var verticesFromCatch = new ICodeUnit[] { catchStatement };

            // If there is a catch expression then build it.
            IEnumerable<ICodeUnit> verticesFromExpression;
            if (catchStatement.CatchExpression != null)
            {
                verticesFromExpression = graph.BuildFlowGraph(
                    verticesFromCatch,
                    catchStatement.CatchExpression);
            }
            else
            {
                verticesFromExpression = verticesFromCatch;
            }

            // Build the graph for the embedded statement.
            var embeddedVertices = graph.BuildFlowGraph(
                verticesFromExpression,
                catchStatement.EmbeddedStatement);

            // The paths out of the catch are:
            //      1) Skip the whole catch.
            //      2) Through the expression and through the embedded statement.
            //      3) Through the expression and skipping the embedded statement.
            IEnumerable<ICodeUnit> lastVertices = verticesFromCatch;
            if (catchStatement.CatchExpression != null)
            {
                lastVertices = lastVertices.Concat(verticesFromExpression);
            }

            return lastVertices.Concat(embeddedVertices);
        }

        private static IEnumerable<ICodeUnit> BuildDoWhileFlowGraph(CodeUnitGraph graph, DoWhileStatement doStatement)
        {
            var verticesFromEmbeddedStatement = graph.BuildFlowGraph(
                new ICodeUnit[] { doStatement },
                doStatement.EmbeddedStatement);

            var verticesFromCondition = graph.BuildFlowGraph(
                verticesFromEmbeddedStatement,
                doStatement.ConditionalExpression);

            // Hook all of the vertices from the condition expression back to the embedded statement to make the loop.
            foreach (var conditionVertex in verticesFromCondition)
            {
                graph.AddEdge(new Edge<ICodeUnit>(conditionVertex, doStatement.EmbeddedStatement));
            }

            return verticesFromCondition;
        }

        private static IEnumerable<ICodeUnit> BuildElseFlowGraph(CodeUnitGraph graph, ElseStatement elseStatement)
        {
            // This is used three times so use the same instance.
            var verticesFromElse = new ICodeUnit[] { elseStatement };

            // If there is a condition then build the condtion.
            IEnumerable<ICodeUnit> verticesFromCondition;
            if (elseStatement.ConditionExpression != null)
            {
                verticesFromCondition = graph.BuildFlowGraph(
                    verticesFromElse,
                    elseStatement.ConditionExpression);
            }
            else
            {
                verticesFromCondition = verticesFromElse;
            }

            // Build the graph for the embedded statement.
            var embeddedVertices = graph.BuildFlowGraph(
                verticesFromCondition,
                elseStatement.EmbeddedStatement);

            // The paths out of the else are:
            //      1) Skip the whole else.
            //      2) Through the condition and then through the embedded statement.
            //      3) Through the condition and skipping the embedded statement.
            IEnumerable<ICodeUnit> lastVertices = verticesFromElse;
            if (elseStatement.ConditionExpression != null)
            {
                lastVertices = lastVertices.Concat(new ICodeUnit[] { elseStatement.ConditionExpression });
            }

            return lastVertices.Concat(embeddedVertices);
        }

        private static IEnumerable<ICodeUnit> BuildFlowGraph(
            this CodeUnitGraph graph,
            IEnumerable<ICodeUnit> lastVertices,
            ICodeUnit codeUnit)
        {
            // Hook all the last vertexes to this item.
            foreach (ICodeUnit lastVertex in lastVertices)
            {
                graph.AddVerticesAndEdge(new Edge<ICodeUnit>(lastVertex, codeUnit));
            }

            // Try each builder. If one returns non-null then the builder did it's job and we can return immediately.
            foreach (CodeUnitBuilder builder in FlowGraph.builders)
            {
                var returnValue = builder(graph, codeUnit);
                if (returnValue != null)
                {
                    return returnValue;
                }
            }

            return graph.BuildFlowGraph(
                new ICodeUnit[] { codeUnit },
                codeUnit.Children());
        }

        private static IEnumerable<ICodeUnit> BuildFlowGraph(
            this CodeUnitGraph graph,
            IEnumerable<ICodeUnit> lastVertices,
            IEnumerable<ICodeUnit> codeUnits)
        {
            foreach (ICodeUnit codeUnit in codeUnits)
            {
                lastVertices = graph.BuildFlowGraph(lastVertices, codeUnit);
            }

            return lastVertices;
        }

        private static IEnumerable<ICodeUnit> BuildForeachFlowGraph(CodeUnitGraph graph, ForeachStatement foreachStatement)
        {
            // Build the variable declaration part of the foreach. This is the indexer part.
            var verticesFromForeach = graph.BuildFlowGraph(
                new ICodeUnit[] { foreachStatement },
                foreachStatement.Variable);

            // Build the enumerable part.
            var verticesFromEnumerablePart = graph.BuildFlowGraph(
                verticesFromForeach,
                foreachStatement.Item);

            // Build the body of the foreach.
            var verticesFromBody = graph.BuildFlowGraph(
                verticesFromEnumerablePart,
                foreachStatement.EmbeddedStatement);

            // The foreach body needs to loop back to its beginning since its a loop.
            foreach (var vertexFromBody in verticesFromBody)
            {
                graph.AddEdge(new Edge<ICodeUnit>(vertexFromBody, foreachStatement.EmbeddedStatement));
            }

            // Since the foreach may never be entered, we have to add the path the goes from the enumerable part.
            return verticesFromEnumerablePart.Concat(verticesFromBody);
        }

        private static IEnumerable<ICodeUnit> BuildForFlowGraph(CodeUnitGraph graph, ForStatement forStatement)
        {
            var verticesFromInitializers = graph.BuildFlowGraph(
                new ICodeUnit[] { forStatement },
                forStatement.Initializers.Cast<ICodeUnit>());

            var verticesFromCondition = graph.BuildFlowGraph(
                verticesFromInitializers,
                forStatement.Condition);

            var verticesFromStatements = graph.BuildFlowGraph(
                verticesFromCondition,
                forStatement.EmbeddedStatement);

            var verticesFromIterators = graph.BuildFlowGraph(
                verticesFromStatements,
                forStatement.Iterators.Cast<ICodeUnit>());

            // Hook the vertexes coming from the iterators to the condition.
            foreach (ICodeUnit vertexFromIterators in verticesFromIterators)
            {
                graph.AddEdge(new Edge<ICodeUnit>(vertexFromIterators, forStatement.Condition));
            }

            return verticesFromCondition;
        }

        private static IEnumerable<ICodeUnit> BuildIfFlowGraph(CodeUnitGraph graph, IfStatement ifStatement)
        {
            var verticesFromCondition = graph.BuildFlowGraph(
                new ICodeUnit[] { ifStatement },
                ifStatement.ConditionExpression);

            var verticesFromTruePart = graph.BuildFlowGraph(
                verticesFromCondition,
                ifStatement.EmbeddedStatement);

            // The vertices from the if statement come from the condition (which is followed if the condition is false)
            // and the embedded statement if the condition is true.
            return verticesFromCondition.Concat(verticesFromTruePart);
        }

        private static IEnumerable<ICodeUnit> BuildLambdaExpressionGraph(CodeUnitGraph graph, LambdaExpression lambdaExpression)
        {
            var verticesFromLambda = graph.BuildFlowGraph(
                new ICodeUnit[] { lambdaExpression },
                lambdaExpression.AnonymousFunctionBody);

            return verticesFromLambda.Concat(new ICodeUnit[] { lambdaExpression });
        }

        private static IEnumerable<ICodeUnit> BuildReturnFlowGraph(CodeUnitGraph graph, ReturnStatement returnStatement)
        {
            // If there is a return value then build it's graph.
            if (returnStatement.ReturnValue != null)
            {
                graph.BuildFlowGraph(
                    new ICodeUnit[] { returnStatement },
                    returnStatement.ReturnValue);
            }

            // The return statement is always the end of the code path so return no last vertices.
            return Enumerable.Empty<ICodeUnit>();
        }

        private static IEnumerable<ICodeUnit> BuildSwitchFlowGraph(CodeUnitGraph graph, SwitchStatement switchStatement)
        {
            // All the switch cases follow the graph for the swith item.
            var verticesFromSwitchItem = graph.BuildFlowGraph(
                new ICodeUnit[] { switchStatement },
                switchStatement.SwitchItem);

            IEnumerable<ICodeUnit> lastVertices;
            if (switchStatement.DefaultStatement == null)
            {
                // If there is no default statement then the code path that skips all the cases is added.
                lastVertices = verticesFromSwitchItem;
            }
            else
            {
                // If there is a default statement then it needs to be in the list of last vertices.
                lastVertices = graph.BuildFlowGraph(
                    verticesFromSwitchItem,
                    switchStatement.DefaultStatement.ChildStatements.Cast<ICodeUnit>());
            }

            // Process each case statement, if any.
            foreach (SwitchCaseStatement caseStatement in switchStatement.CaseStatements)
            {
                // Graph the swith item's identifier.
                var caseLastVertices = graph.BuildFlowGraph(
                    verticesFromSwitchItem,
                    caseStatement.Identifier);

                // Graph all the statements in the case.
                caseLastVertices = graph.BuildFlowGraph(
                    caseLastVertices,
                    caseStatement.ChildStatements.Cast<ICodeUnit>());

                // Concatenate these last vertices with the rest.
                lastVertices = lastVertices.Concat(caseLastVertices);
            }

            return lastVertices;
        }

        private static IEnumerable<ICodeUnit> BuildThrowStatementGraph(CodeUnitGraph graph, ThrowStatement throwStatement)
        {
            // Build the throw expression.
            if (throwStatement.ThrownExpression != null)
            {
                graph.BuildFlowGraph(
                    new ICodeUnit[] { throwStatement },
                    throwStatement.ThrownExpression);
            }

            // The throw statement is always the end of the code path so return no last vertices.
            return Enumerable.Empty<ICodeUnit>();
        }

        private static IEnumerable<ICodeUnit> BuildTryFlowGraph(CodeUnitGraph graph, TryStatement tryStatement)
        {
            // The try statement is the last vertices for the all the different parts so just use one instance.
            var tryStatementVertices = new ICodeUnit[] { tryStatement };

            var verticesFromEmbedded = graph.BuildFlowGraph(tryStatementVertices, tryStatement.EmbeddedStatement);

            // There are two paths out of the try.
            //  1) Bypass the try embedded statement to simulate an exception at the very beginning.
            //  2) Through the embedded statement.
            return tryStatementVertices.Concat(verticesFromEmbedded);
        }

        private static IEnumerable<ICodeUnit> BuildWhileFlowGraph(CodeUnitGraph graph, WhileStatement whileStatement)
        {
            var verticesFromCondition = graph.BuildFlowGraph(
                new ICodeUnit[] { whileStatement },
                whileStatement.ConditionExpression);

            var verticesFromBody = graph.BuildFlowGraph(
                verticesFromCondition,
                whileStatement.EmbeddedStatement);

            // Hook the vertexes coming from the embedded statement to the condition.
            foreach (ICodeUnit vertexFromIterators in verticesFromBody)
            {
                graph.AddEdge(new Edge<ICodeUnit>(vertexFromIterators, whileStatement.ConditionExpression));
            }

            // The code paths are as follows:
            //      1) Through the condition and the body. The body wraps back to the condition.
            //      2) Through the condition and skip the body.
            return verticesFromCondition.Concat(verticesFromBody);
        }

        private static CodeUnitBuilder CreateBuilder<TCode>(Func<CodeUnitGraph, TCode, IEnumerable<ICodeUnit>> builder)
            where TCode : class
        {
            return
                (graph, codeUnit) =>
                {
                    TCode typedCodeUnit = codeUnit as TCode;
                    if (typedCodeUnit != null)
                    {
                        return builder(graph, typedCodeUnit);
                    }
                    else
                    {
                        return null;
                    }
                };
        }
        #endregion Private Methods
    }
}